import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch
from matplotlib.collections import LineCollection
from matplotlib import cm

# 设置全局字体

plt.rcParams.update({
    'font.size': 10,          # 全局字体大小
    'axes.titlesize': 16,     # 标题字体大小
    'axes.labelsize': 10,     # 坐标轴标签字体大小
    'xtick.labelsize': 10,    # x轴刻度字体
    'ytick.labelsize': 10,    # y轴刻度字体
    'legend.fontsize': 10     # 图例字体
})
# 初始化设置
plt.rcParams['font.sans-serif'] = ['KaiTi']  # 中文字体
plt.rcParams['mathtext.fontset'] = 'stix'     # 数学字体
plt.rcParams['axes.unicode_minus'] = False




# 定义目标函数
def f(x, y):
    return 0.5 * x**2 + 2.5 * y**2

# 定义梯度函数
def gradient(x, y):
    grad_x = x  # 对x的偏导数
    grad_y = 5 * y  # 对y的偏导数
    return (grad_x, grad_y)

# 设置起点
x_start, y_start = 3.0, 3.0

# 设置学习率和迭代次数
learning_rate = 0.1
steps = 15

# 存储路径
path = [(x_start, y_start)]
x_values, y_values = [x_start], [y_start]

# 执行梯度下降
current_x, current_y = x_start, y_start
for i in range(steps):
    grad_x, grad_y = gradient(current_x, current_y)
    # 更新位置
    new_x = current_x - learning_rate * grad_x
    new_y = current_y - learning_rate * grad_y
    path.append((new_x, new_y))
    x_values.append(new_x)
    y_values.append(new_y)
    current_x, current_y = new_x, new_y

# 转换为numpy数组便于绘图
path = np.array(path)
x_values = np.array(x_values)
y_values = np.array(y_values)

# 创建网格用于绘制函数等高线
x_grid = np.linspace(-5, 5, 100)
y_grid = np.linspace(-5, 5, 100)
X, Y = np.meshgrid(x_grid, y_grid)
Z = f(X, Y)

# 绘制图形
plt.figure(figsize=(10, 8))
contour = plt.contour(X, Y, Z, levels=20, cmap='viridis')
# plt.colorbar(contour, label='函数值')

# 绘制梯度下降路径
plt.plot(x_values, y_values, 'ro-', linewidth=2, markersize=6, label='梯度下降路径')

# 标记起点和终点
plt.scatter(x_start, y_start, c='green', s=100, marker='o', label='起点')
plt.scatter(0, 0, c='blue', s=100, marker='*', label='最小值点')

# 绘制每一步的梯度方向

for i in range(len(path) - 1):
    x, y = path[i]
    grad_x, grad_y = gradient(x, y)
    # 绘制负梯度方向（梯度下降方向）
    dx = -grad_x * 0.5  # 缩放系数，使箭头更明显
    dy = -grad_y * 0.5
    plt.arrow(x, y, dx, dy, head_width=0.1, head_length=0.2, fc='purple', ec='purple', label='梯度方向' if i == 0 else "")

# 添加标签和图例
plt.xlabel('x')
plt.ylabel('y')
plt.title('函数 $f(x,y) = 0.5x^2 + 2.5y^2$ 的梯度下降法演示')
plt.legend()
plt.grid(True)
plt.axis('equal')
# plt.tight_layout()
plt.show()